# -*- coding: utf-8 -*-

import torch
from transform import target_transform


class TCSMLoss(torch.nn.Module):
    def __init__(self, max_mse_weight = 1.0, device='cuda'):
        super(TCSMLoss, self).__init__()
        self.max_mse_weight = max_mse_weight
        self.device = device
        self.__name__ = "TCSMLoss"
    
    def forward(
            self,
            labeled_student_pred,
            labeled_teacher_pred,
            unlabeled_student_pred,
            unlabeled_teacher_pred,
            target,
            labeled_trans_args,
            unlabeled_trans_args,
            epoch
        ):
        labeled_student_pred = labeled_student_pred.to(self.device)
        labeled_teacher_pred = labeled_teacher_pred.to(self.device)
        unlabeled_student_pred = unlabeled_student_pred.to(self.device)
        unlabeled_teacher_pred = unlabeled_teacher_pred.to(self.device)
        target = target.to(self.device).long()
        
        target = target_transform(target, labeled_trans_args)
        
        labeled_student_pred = target_transform(labeled_student_pred, labeled_trans_args)
        unlabeled_student_pred = target_transform(unlabeled_student_pred, unlabeled_trans_args)
        
        cross_entropy = torch.nn.CrossEntropyLoss()(labeled_student_pred, target[:,0,...])        
        labeled_mse = torch.nn.MSELoss()(labeled_student_pred, labeled_teacher_pred)
        unlabeled_mse = torch.nn.MSELoss()(unlabeled_student_pred, unlabeled_teacher_pred)
        
        labeled_batch_size = labeled_student_pred.shape[0]
        unlabeled_batch_size = unlabeled_student_pred.shape[0]
        mse = (labeled_batch_size * labeled_mse + unlabeled_batch_size * unlabeled_mse) / (labeled_batch_size + unlabeled_batch_size)
        
        mse_weight = self.max_mse_weight * gaussian_rampup(epoch)
        tot_loss = cross_entropy + mse_weight * mse
        
        return tot_loss


class CrossEntropyLoss(torch.nn.Module):
    def __init__(self, device='cuda'):
        super(CrossEntropyLoss, self).__init__()
        self.__name__ = "CrossEntropyLoss"
    
    def forward(self, y, y_pred):
        y_pred = y_pred[:,0,...]
        return torch.nn.CrossEntropyLoss()(y, y_pred)


def gaussian_rampup(epoch, rampup_length=400):
    epoch = torch.tensor(epoch)
    epoch = torch.clamp(epoch, min=0, max=rampup_length)
    phase = 1.0 - epoch / rampup_length
    return torch.exp(torch.tensor(-5.0) * torch.pow(phase, 2))


# if __name__ == "__main__":
#     from CC_CCI_dataset import get_TCSM_dataloader
#     from model import TCSM
#     import segmentation_models_pytorch as smp
#     # import matplotlib.pyplot as plt
#     IMAGE_DIR = "D:\data\ct_lesion_seg\image"
#     MASK_DIR = "D:\data\ct_lesion_seg\mask"
#     train_loader, valid_loader, unlabeled_loader = get_TCSM_dataloader(IMAGE_DIR, MASK_DIR, 5)
    
#     student = smp.DeepLabV3Plus(encoder_weights=None, in_channels=1, classes=3, activation='softmax2d')
#     teacher = smp.DeepLabV3Plus(encoder_weights=None, in_channels=1, classes=3, activation='softmax2d')
#     tcsm = TCSM(student, teacher)
    
#     train_loader_iter = iter(train_loader)
#     unlabeled_loader_iter = iter(unlabeled_loader)
#     for i in range(1):
#         x, y = train_loader_iter.next()
#         z = unlabeled_loader_iter.next()
        
#         ls, lt, largs = tcsm(x)
#         # student(x.cuda())
#         us, ut, uargs = tcsm(y)
    
#     print("breakpoint")
#     loss = TCSMLoss()(ls, lt, us, ut, y, largs, uargs, 0)
#     loss.backward()
    